package cn.kk.xss;

import com.alibaba.fastjson.JSON;
import jakarta.servlet.ReadListener;
import jakarta.servlet.ServletInputStream;
import jakarta.servlet.http.HttpServletRequest;
import jakarta.servlet.http.HttpServletRequestWrapper;
import lombok.Getter;
import lombok.extern.slf4j.Slf4j;
import org.apache.commons.lang3.StringUtils;

import java.io.*;
import java.nio.charset.StandardCharsets;
import java.util.HashMap;
import java.util.Map;

/**
 * @author kinoko
 */
@Slf4j
@Getter
public class XssHttpServletRequestWrapper extends HttpServletRequestWrapper {

	HttpServletRequest orgRequest;

	public XssHttpServletRequestWrapper(HttpServletRequest request) {
		super(request);
		orgRequest = request;
	}

	/**
	 * 获取最原始的request的静态方法
	 * @return R
	 */
	public static HttpServletRequest getOrgRequest(HttpServletRequest req) {
		if (req instanceof XssHttpServletRequestWrapper) {
			return ((XssHttpServletRequestWrapper) req).getOrgRequest();
		}
		return req;
	}

	/**
	 * 核心 XSS检测,处理(通过改变这个对象来达到执行不同XSS预防策略实现)
	 * @param str 字符串
	 * @return 处理结果
	 */
	private String xssEncode(String str) {
		// 通过改变这个对象来达到执行不同XSS预防策略实现
		return XssEncode.xssEncode(str);
	}

	/**
	 * 覆盖getParameter方法，将参数名和参数值都做xss过滤。<br/>
	 * 如果需要获得原始的值，则通过super.getParameterValues(name)来获取<br/>
	 * getParameterNames,getParameterValues和getParameterMap也可能需要覆盖
	 */
	@Override
	public String getParameter(String name) {
		String value = super.getParameter(xssEncode(name));
		if (value != null) {
			value = xssEncode(value);
		}
		return value;
	}

	/**
	 * 重写getParameterMap
	 */
	@Override
	public Map<String, String[]> getParameterMap() {
		HashMap<String, String[]> paramMap = (HashMap<String,
			String[]>) super.getParameterMap();
		paramMap = (HashMap<String, String[]>) paramMap.clone();
		for (Map.Entry<String, String[]> entry : paramMap.entrySet()) {
			String[] values = entry.getValue();
			for (int i = 0; i < values.length; i++) {
				if (values[i] != null) {
					values[i] = xssEncode(values[i]);
				}
			}
			entry.setValue(values);
		}
		return paramMap;
	}

	/**
	 * 重写getParameterValues
	 */
	@Override
	public String[] getParameterValues(String name) {
		String[] values = super.getParameterValues(name);
		if (values == null) {
			values = new String[]{};
		}
		int count = values.length;
		String[] encodedValues = new String[count];
		for (int i = 0; i < count; i++) {
			encodedValues[i] = xssEncode(values[i]);
		}
		return encodedValues;
	}

	/**
	 * 覆盖getHeader方法，将参数名和参数值都做xss过滤。<br/>
	 * 如果需要获得原始的值，则通过super.getHeaders(name)来获取<br/>
	 * getHeaderNames 也可能需要覆盖
	 */
	@Override
	public String getHeader(String name) {
		String value = super.getHeader(xssEncode(name));
		if (value != null) {
			value = xssEncode(value);
		}
		return value;
	}

	@Override
	public String getQueryString() {
		return xssEncode(super.getQueryString());
	}

	@Override
	public ServletInputStream getInputStream() throws IOException {
		String str = getRequestBody(super.getInputStream());
		if (StringUtils.isNoneBlank(str)) {
			Map<String, Object> map = JSON.parseObject(str, Map.class);
			Map<String, Object> resultMap = new HashMap<>(map.size());
			for (String key : map.keySet()) {
				Object val = map.get(key);
				if (map.get(key) instanceof String) {
					resultMap.put(key, xssEncode(val.toString()));
				} else {
					resultMap.put(key, val);
				}
			}
			str = JSON.toJSONString(resultMap);
		}
		final ByteArrayInputStream bais = new ByteArrayInputStream(str.getBytes());
		return new ServletInputStream() {
			@Override
			public int read() {
				return bais.read();
			}

			@Override
			public boolean isFinished() {
				return false;
			}

			@Override
			public boolean isReady() {
				return false;
			}

			@Override
			public void setReadListener(ReadListener listener) {
			}
		};
	}

	/**
	 * 处理比较的json字符串
	 * @param stream	输入流
	 * @return	R
	 */
	private String getRequestBody(InputStream stream) {
		String line;
		StringBuilder body = new StringBuilder();
		// 读取POST提交的数据内容
		BufferedReader reader = new BufferedReader(new InputStreamReader(stream, StandardCharsets.UTF_8));
		try {
			while ((line = reader.readLine()) != null) {
				body.append(line);
			}
		} catch (IOException e) {
			log.error("[XSS策略异常]" + e.getMessage());
		}
		return body.toString();
	}

}